#include "KDTree.h"

KDTree::KDTree(size_t _targetTriangles, size_t _maxTreeDepth): targetTriangles(_targetTriangles), maxTreeDepth(_maxTreeDepth){
}

KDTree::~KDTree(void){
}

void KDTree::buildTree(vector<Triangle> *sceneTriangles, BoundingBox sceneBox){
	Vector3 nextNormal = Vector3(1,0,0);
	KDTreeNode* newFront = new KDTreeNode();
	KDTreeNode* newBack  = new KDTreeNode();

	Vector3 boxRange = sceneBox.maxValues - sceneBox.minValues;
	BoundingBox newFrontBox = sceneBox;
	BoundingBox newBackBox  = sceneBox;
	newFrontBox.maxValues.z = sceneBox.maxValues.z-(boxRange.z*0.5);
	newBackBox.minValues.z  = sceneBox.maxValues.z-(boxRange.z*0.5);
	
	root.planeNormal = Vector3(0,0,1);
	root.position = sceneBox.maxValues-(boxRange*0.5);
	root.AABB = sceneBox;

	//Divide triangles in to child boxes.
	bool frontAdded;
	bool backAdded;
	int axis = 2;
	for(int i = 0; i < (int)sceneTriangles->size(); i++){
		frontAdded = backAdded = false;
		for(int j = 0; j < 3; j++){
			if(!frontAdded && sceneTriangles->at(i).verticies.at(j)[axis] <= root.position[axis]){
				newFront->addTriangle(&(sceneTriangles->at(i)));
				frontAdded = true;
			}
			if(!backAdded && sceneTriangles->at(i).verticies.at(j)[axis] >= root.position[axis]){
				newBack->addTriangle(&(sceneTriangles->at(i)));
				backAdded = true;
			}
		}
	}

	newFront->planeNormal = nextNormal;
	newBack->planeNormal  = nextNormal;
	newFront->AABB = newFrontBox;
	newBack->AABB = newBackBox;
	root.frontChild = newFront;
	root.backChild  = newBack;

	buildTree(newFront, 0, 1);
	buildTree(newBack,  0, 1);
}

void KDTree::buildTree(KDTreeNode *curNode, int axis, size_t depth){
	int nextAxis = 0;
	if(axis != 2){
		nextAxis = axis+1;
	}

	//If the current node contains more than the number of triangles desired and is not past the max recursive depth, generate another split.
	if(curNode->triangles.size() > targetTriangles && depth < maxTreeDepth){
		Vector3 nextNormal(0,0,0);
		KDTreeNode* newFront = new KDTreeNode();
		KDTreeNode* newBack  = new KDTreeNode();
		BoundingBox newFrontBox = curNode->AABB;
		BoundingBox newBackBox  = curNode->AABB;
		Vector3 boxRange = curNode->AABB.maxValues - curNode->AABB.minValues;

		curNode->position = curNode->AABB.maxValues-(boxRange*0.5);
		if(axis == 0){
			nextNormal.y = 1;
			newFrontBox.maxValues.x = curNode->AABB.maxValues.x-(boxRange.x*0.5);
			newBackBox.minValues.x  = curNode->AABB.maxValues.x-(boxRange.x*0.5);
		}
		else if(axis == 1){
			nextNormal.z = 1;
			newFrontBox.maxValues.y = curNode->AABB.maxValues.y-(boxRange.y*0.5);
			newBackBox.minValues.y  = curNode->AABB.maxValues.y-(boxRange.y*0.5);
		}
		else if(axis == 2){
			nextNormal.x = 1;
			newFrontBox.maxValues.z = curNode->AABB.maxValues.z-(boxRange.z*0.5);
			newBackBox.minValues.z  = curNode->AABB.maxValues.z-(boxRange.z*0.5);
		}

		//Divide triangles in to child boxes.
		bool frontAdded;
		bool backAdded;
		for(size_t i = 0; i < curNode->triangles.size(); i++){
			frontAdded = backAdded = false;
			for(int j = 0; j < 3; j++){
				if(!frontAdded && curNode->triangles.at(i)->verticies.at(j)[axis] <= curNode->position[axis]){
					newFront->addTriangle(curNode->triangles.at(i));
					frontAdded = true;
				}
				if(!backAdded && curNode->triangles.at(i)->verticies.at(j)[axis] >= curNode->position[axis]){
					newBack->addTriangle(curNode->triangles.at(i));
					backAdded = true;
				}
			}
		}

		newFront->planeNormal = nextNormal;
		newBack->planeNormal  = nextNormal;
		newFront->AABB = newFrontBox;
		newBack->AABB = newBackBox;
		curNode->frontChild = newFront;
		curNode->backChild  = newBack;

		buildTree(newFront, nextAxis, depth+1);
		buildTree(newBack,  nextAxis, depth+1);
	}
}

vector<Triangle*> KDTree::getPossibleIntersections(Ray r, Point3 boxIn, Point3 boxOut){
	return getPossibleIntersections(r, &root, boxIn, boxOut, 2);
}

vector<Triangle*> KDTree::getPossibleIntersections(Ray r, KDTreeNode* curNode, Point3 boxIn, Point3 boxOut, int axis){
	r.direction.normalize();
	//If this is a leaf node, return all triangles associated with this box.
	if(curNode->frontChild == NULL && curNode->backChild == NULL){
		return curNode->triangles;
	}
	//Otherwise traverse the tree
	else{
		int nextAxis = 0;
		if(axis != 2){
			nextAxis = axis+1;
		}

		//Check if ray is parallel to curNode's plane
		if(r.direction*curNode->planeNormal == 0){
			//Check if ray lies on the plane (traverse both sides)
			if(r.origin[axis] == curNode->position[axis]){
				vector<Triangle*> frontObjects = getPossibleIntersections(r, curNode->frontChild, curNode->position, curNode->position, nextAxis);
				vector<Triangle*> backObjects = getPossibleIntersections(r, curNode->backChild, curNode->position, curNode->position, nextAxis);
				for(int i = 0; i < (int)backObjects.size(); i++){
					frontObjects.push_back(backObjects.at(i));
				}
				return frontObjects;
			}
			//Check if ray is in front of the line (traverse front branch)
			else if(r.origin[axis] < curNode->position[axis]){
				return getPossibleIntersections(r, curNode->frontChild, boxIn, boxOut, nextAxis);
			}
			//Otherwise, ray is in back of the line (traverse back branch)
			else{
				return getPossibleIntersections(r, curNode->backChild, boxIn, boxOut, nextAxis);
			}
		}
		else{
			double u = (curNode->planeNormal*(curNode->position-r.origin))/(r.direction*curNode->planeNormal);
			Point3 planeIntersection = r.origin+(u*r.direction);
			if(boxIn[axis] <= planeIntersection[axis]){
				if(boxOut[axis] < planeIntersection[axis]){
					return getPossibleIntersections(r, curNode->frontChild, boxIn, boxOut, nextAxis);
				}
				/*else if(boxOut[axis] == planeIntersection[axis]){
				}*/
				else{
					vector<Triangle*> backObjects = getPossibleIntersections(r, curNode->frontChild, boxIn, planeIntersection, nextAxis);
					vector<Triangle*> frontObjects = getPossibleIntersections(r, curNode->backChild, planeIntersection, boxOut, nextAxis);
					for(int i = 0; i < (int)frontObjects.size(); i++){
						backObjects.push_back(frontObjects.at(i));
					}
					return backObjects;
				}
			}
			else{
				if(boxOut[axis] > planeIntersection[axis]){
					return getPossibleIntersections(r, curNode->backChild, boxIn, boxOut, nextAxis);
				}
				else{
					vector<Triangle*> frontObjects = getPossibleIntersections(r, curNode->backChild, boxIn, planeIntersection, nextAxis);
					vector<Triangle*> backObjects = getPossibleIntersections(r, curNode->frontChild, planeIntersection, boxOut, nextAxis);
					for(int i = 0; i < (int)backObjects.size(); i++){
						frontObjects.push_back(backObjects.at(i));
					}
					return frontObjects;
				}
			}
		}
	}
}